package com.wuwei.elearning.question.template;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Scanner;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import java.lang.System;

public class AiAsk {
    private static final Logger logger = LoggerFactory.getLogger(AiAsk.class);
    private static final List<Message> conversationHistory = new ArrayList<>();

    private static void handleStreamingOutput(GenerationResult message) {
        if (message.getOutput() != null && message.getOutput().getChoices() != null &&
                !message.getOutput().getChoices().isEmpty() &&
                message.getOutput().getChoices().get(0).getMessage() != null) {

            String content = message.getOutput().getChoices().get(0).getMessage().getContent();
            if (content != null) {
                System.out.print(content);
            }
        }
    }

    private static Message createAssistantMessage(String content) {
        return Message.builder().role(Role.ASSISTANT.getValue()).content(content).build();
    }

    private static Message createUserMessage(String content) {
        return Message.builder().role(Role.USER.getValue()).content(content).build();
    }

    public static void streamCallWithMessageHistory(Generation gen, String userInput)
            throws NoApiKeyException, ApiException, InputRequiredException {
        // Add user message to conversation history
        Message userMsg = createUserMessage(userInput);
        conversationHistory.add(userMsg);

        // Create parameter with conversation history
        GenerationParam param = buildGenerationParam();

        // Call API with streaming output
        StringBuilder fullResponse = new StringBuilder();
        Flowable<GenerationResult> result = gen.streamCall(param);
        result.blockingForEach(message -> {
            handleStreamingOutput(message);
            if (message.getOutput() != null && message.getOutput().getChoices() != null &&
                    !message.getOutput().getChoices().isEmpty() &&
                    message.getOutput().getChoices().get(0).getMessage() != null) {

                String content = message.getOutput().getChoices().get(0).getMessage().getContent();
                if (content != null) {
                    fullResponse.append(content);
                }
            }
        });

        // Add assistant's response to conversation history
        conversationHistory.add(createAssistantMessage(fullResponse.toString()));
        System.out.println("\n");
    }

    private static GenerationParam buildGenerationParam() {
        String dashscopeApiKey = System.getenv("DASHSCOPE_API_KEY");
        if (dashscopeApiKey == null || dashscopeApiKey.isEmpty()) {
            logger.warn("DASHSCOPE_API_KEY environment variable not set. Please set it or provide API key directly.");
            // You could replace this with your own API key if needed
            // dashscopeApiKey = "sk-your-api-key";
        }

        return GenerationParam.builder()
                .apiKey(dashscopeApiKey)
                // Model is qwen-plus, can be changed to other models
                .model("qwen-plus")
                .messages(conversationHistory)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        Generation gen = new Generation();

        // Add system message at the beginning (optional)
        conversationHistory.add(Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("你是阿里云智能助手，请提供有帮助的回答。")
                .build());

        System.out.println("欢迎使用阿里云百炼对话系统，输入'退出'结束对话");

        while (true) {
            System.out.print("用户: ");
            String userInput = scanner.nextLine().trim();

            if (userInput.equalsIgnoreCase("退出")) {
                System.out.println("谢谢使用，再见！");
                break;
            }

            System.out.print("助手: ");
            try {
                streamCallWithMessageHistory(gen, userInput);
            } catch (ApiException | NoApiKeyException | InputRequiredException e) {
                logger.error("调用API时发生错误: {}", e.getMessage());
                System.out.println("抱歉，调用API时发生错误: " + e.getMessage());
            }
        }

        scanner.close();
    }
}
